{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import torch\n",
    "import pandas as pd\n",
    "import numpy as np \n",
    "import xml.etree.ElementTree as ET\n",
    "from scipy.io import loadmat\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from PIL import Image\n",
    "from tqdm.notebook import tqdm\n",
    "from sklearn import metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "def load_dataset():\n",
    "\n",
    "    # Get train list\n",
    "    f = loadmat(\"lists/train_list.mat\")\n",
    "    train_images = [x[0][0] for x in f['file_list']]\n",
    "    train_labels = [x[0] for x in f['labels']]\n",
    "\n",
    "    # Get file list\n",
    "    f = loadmat(\"lists/test_list.mat\")\n",
    "    test_images = [x[0][0] for x in f['file_list']]\n",
    "    test_labels = [x[0] for x in f['labels']]\n",
    "\n",
    "    # Gather data\n",
    "    train_data = []\n",
    "    test_data = []\n",
    "\n",
    "    # Record category ids\n",
    "    categories = {}\n",
    "\n",
    "    for i in range(len(train_images) + len(test_images)):\n",
    "\n",
    "        # Determine if train or test\n",
    "        image = train_images[i] if i < len(train_images) else test_images[i - len(train_images)]\n",
    "        label = train_labels[i] if i < len(train_images) else test_labels[i - len(train_images)]\n",
    "        label_name = os.path.split(image)[0]\n",
    "        # Label -1 to make it 0-indexed\n",
    "        categories[label_name] = label-1\n",
    "        annotation_path = os.path.join(\"Annotation\", image.replace(\".jpg\", \"\"))\n",
    "\n",
    "        # Read XML\n",
    "        tree = ET.parse(annotation_path)\n",
    "        root = tree.getroot()\n",
    "\n",
    "        width = int(root.find(\"size\").find(\"width\").text)\n",
    "        height = int(root.find(\"size\").find(\"height\").text)\n",
    "\n",
    "        bndbox = root.find(\"object\").find(\"bndbox\")\n",
    "        xmin = int(bndbox.find(\"xmin\").text)\n",
    "        ymin = int(bndbox.find(\"ymin\").text)\n",
    "        xmax = int(bndbox.find(\"xmax\").text)\n",
    "        ymax = int(bndbox.find(\"ymax\").text)\n",
    "\n",
    "        # Append to data\n",
    "        if i < len(train_images):\n",
    "            train_data.append(dict(\n",
    "                image=os.path.join(\"Images\", image),\n",
    "                label=label-1,\n",
    "                label_name=label_name,\n",
    "                width=width,\n",
    "                height=height,\n",
    "                xmin=xmin,\n",
    "                ymin=ymin,\n",
    "                xmax=xmax,\n",
    "                ymax=ymax\n",
    "            ))\n",
    "        else:\n",
    "            test_data.append(dict(\n",
    "                image=os.path.join(\"Images\", image),\n",
    "                label=label-1,\n",
    "                label_name=label_name,\n",
    "                width=width,\n",
    "                height=height,\n",
    "                xmin=xmin,\n",
    "                ymin=ymin,\n",
    "                xmax=xmax,\n",
    "                ymax=ymax\n",
    "            ))\n",
    "\n",
    "\n",
    "    return train_data, test_data, categories\n",
    "\n",
    "# Read dataset and gather into dataframe\n",
    "train_data, test_data, categories = load_dataset()\n",
    "dftrain = pd.DataFrame(train_data)\n",
    "dftest = pd.DataFrame(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of classes:  120\n",
      "Number of training samples:  12000\n",
      "Number of testing samples:  8580\n"
     ]
    }
   ],
   "source": [
    "# Get the classes summary\n",
    "print(\"Number of classes: \", len(categories))\n",
    "print(\"Number of training samples: \", len(dftrain))\n",
    "print(\"Number of testing samples: \", len(dftest))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[[-1.5699, -1.6384, -1.6727,  ..., -1.3644, -1.3644, -1.4158],\n",
       "          [-1.4500, -1.5699, -1.6213,  ..., -1.3644, -1.3644, -1.4158],\n",
       "          [-1.3644, -1.4843, -1.5528,  ..., -1.3644, -1.3644, -1.4158],\n",
       "          ...,\n",
       "          [ 0.3823,  0.3138,  0.3481,  ..., -0.4568, -0.5082, -0.2513],\n",
       "          [ 0.3823,  0.3138,  0.3481,  ..., -0.4739, -0.5424, -0.2513],\n",
       "          [ 0.3652,  0.3138,  0.3481,  ..., -0.3541, -0.4397, -0.3883]],\n",
       " \n",
       "         [[-1.6506, -1.5980, -1.5455,  ..., -1.3704, -1.3704, -1.4230],\n",
       "          [-1.5980, -1.5630, -1.5105,  ..., -1.3704, -1.3704, -1.4230],\n",
       "          [-1.5455, -1.5455, -1.4930,  ..., -1.3704, -1.3704, -1.4230],\n",
       "          ...,\n",
       "          [-0.0924, -0.1450, -0.1099,  ...,  0.2227,  0.1702,  0.4328],\n",
       "          [-0.0924, -0.1450, -0.1099,  ...,  0.2052,  0.1352,  0.4328],\n",
       "          [-0.0924, -0.1450, -0.1099,  ...,  0.2927,  0.2052,  0.2577]],\n",
       " \n",
       "         [[-1.6302, -1.5779, -1.4733,  ..., -1.5604, -1.5604, -1.6127],\n",
       "          [-1.5953, -1.5430, -1.4733,  ..., -1.5604, -1.5604, -1.6127],\n",
       "          [-1.5779, -1.5430, -1.4733,  ..., -1.5604, -1.5604, -1.6127],\n",
       "          ...,\n",
       "          [-0.5321, -0.5844, -0.5495,  ...,  0.2522,  0.1999,  0.4614],\n",
       "          [-0.5321, -0.5844, -0.5495,  ...,  0.2348,  0.1651,  0.4614],\n",
       "          [-0.5321, -0.5670, -0.5321,  ...,  0.3393,  0.2522,  0.3045]]]),\n",
       " 0)"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Inherit from Dataset\n",
    "class CustomDataset(Dataset):\n",
    "\n",
    "    def __init__(self, df, transform=None):\n",
    "        self.df = df\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        row = self.df.iloc[idx]\n",
    "        image = Image.open(row['image'])\n",
    "        image = image.convert('RGB')\n",
    "        label = row['label']\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        return image, label\n",
    "        \n",
    "\n",
    "train_transforms = transforms.Compose([\n",
    "    # Randomly resize and crop the image to 224\n",
    "    transforms.RandomResizedCrop(224),\n",
    "    # Randomly flip the image horizontally\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    # Convert the image to a PyTorch Tensor\n",
    "    transforms.ToTensor(),\n",
    "    # Normalize the image\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "test_transforms = transforms.Compose([\n",
    "    # Resize the image to 256\n",
    "    transforms.Resize(256),\n",
    "    # Crop the center of the image\n",
    "    transforms.CenterCrop(224),\n",
    "    # Convert the image to a PyTorch Tensor\n",
    "    transforms.ToTensor(),\n",
    "    # Normalize the image\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "\n",
    "\n",
    "train_dataset = CustomDataset(dftrain, transform=train_transforms)\n",
    "test_dataset = CustomDataset(dftest, transform=test_transforms)\n",
    "\n",
    "# Test\n",
    "train_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3854d32961ac4ab8963b9cbcd318e28f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 0:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bb21822aee9b42c4b12ed36ae92b5229",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 0:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  0\n",
      "Accuracy:  0.6665501165501165\n",
      "F1 Score:  0.6416622809566006\n",
      "Confusion Matrix:  [[ 28   0   0 ...   1   0   0]\n",
      " [  0  74   0 ...   0   0   0]\n",
      " [  0   0 112 ...   0   0   0]\n",
      " ...\n",
      " [  1   0   0 ...  41   6   0]\n",
      " [  0   0   0 ...   1  44   0]\n",
      " [  0   0   0 ...   0   1  63]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "97738e053b4349ec9078877fefeae474",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 1:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "79aa487ba15e47a7891e022a6ec1b857",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 1:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  1\n",
      "Accuracy:  0.7092074592074592\n",
      "F1 Score:  0.6947704937888362\n",
      "Confusion Matrix:  [[ 33   0   0 ...   1   0   0]\n",
      " [  0  74   0 ...   0   0   0]\n",
      " [  0   0 132 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  40  10   0]\n",
      " [  0   0   0 ...   1  46   1]\n",
      " [  0   0   0 ...   0   2  61]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "26ad000873234232816793b8306895f0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 2:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eb3a45406ef14e22920a513b2f33cd77",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 2:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  2\n",
      "Accuracy:  0.7403263403263404\n",
      "F1 Score:  0.728980014995554\n",
      "Confusion Matrix:  [[ 39   0   0 ...   0   0   0]\n",
      " [  0  75   0 ...   0   0   0]\n",
      " [  0   0 123 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  43   3   0]\n",
      " [  0   0   0 ...   1  47   0]\n",
      " [  0   0   0 ...   0   1  58]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a58843b3a164308a4b0c3a4df2c6028",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 3:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e686a5e9d81840b6b815e720978f90f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 3:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  3\n",
      "Accuracy:  0.7452214452214452\n",
      "F1 Score:  0.7332687400667351\n",
      "Confusion Matrix:  [[ 36   0   0 ...   1   0   0]\n",
      " [  0  78   0 ...   0   0   0]\n",
      " [  0   0 118 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  42   7   0]\n",
      " [  0   0   0 ...   1  48   1]\n",
      " [  0   0   0 ...   0   2  64]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4860079b86ef45db8e9693145a39da6e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 4:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "796ba7ef469e4e8b80a8569a302b88db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 4:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  4\n",
      "Accuracy:  0.7407925407925408\n",
      "F1 Score:  0.7306036438579546\n",
      "Confusion Matrix:  [[ 37   0   0 ...   2   0   0]\n",
      " [  0  81   0 ...   0   0   0]\n",
      " [  0   0 105 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  51   0   0]\n",
      " [  0   0   0 ...   3  43   1]\n",
      " [  0   0   0 ...   0   2  65]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "12d1528777d546ec9a34a9f806567293",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 5:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a685c4800019485687e858a2c3663367",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 5:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  5\n",
      "Accuracy:  0.7508158508158508\n",
      "F1 Score:  0.7383518863885412\n",
      "Confusion Matrix:  [[ 36   0   0 ...   0   0   0]\n",
      " [  0  74   0 ...   0   0   0]\n",
      " [  0   0 135 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  44   5   0]\n",
      " [  0   0   0 ...   1  47   1]\n",
      " [  0   0   0 ...   0   2  60]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "167533149cc145beb27a8e107b611d46",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 6:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a5fba3a715e450197ef334c04e00e18",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 6:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  6\n",
      "Accuracy:  0.7458041958041958\n",
      "F1 Score:  0.7355986792846984\n",
      "Confusion Matrix:  [[ 35   0   0 ...   0   0   0]\n",
      " [  0  73   0 ...   0   0   0]\n",
      " [  0   0 128 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  33  12   0]\n",
      " [  0   0   0 ...   0  48   0]\n",
      " [  0   0   0 ...   0   4  60]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7573fb33092749d1ad32179d53f6bdd1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 7:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1905a45f3ddc4ce9a0e76097014453f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 7:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  7\n",
      "Accuracy:  0.7353146853146854\n",
      "F1 Score:  0.7255668359090135\n",
      "Confusion Matrix:  [[ 39   0   0 ...   0   0   0]\n",
      " [  0  75   0 ...   0   0   0]\n",
      " [  0   0 126 ...   0   0   0]\n",
      " ...\n",
      " [  1   0   0 ...  34  10   0]\n",
      " [  0   0   0 ...   1  47   1]\n",
      " [  0   0   0 ...   0   3  61]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc0a65ac8baa458692a5bbcf181e9bea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 8:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c1c46e09e54b44e9bf8a037fd3635509",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 8:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  8\n",
      "Accuracy:  0.7567599067599068\n",
      "F1 Score:  0.7484446624604898\n",
      "Confusion Matrix:  [[ 40   0   0 ...   0   0   0]\n",
      " [  0  78   0 ...   0   0   0]\n",
      " [  0   0 126 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  48   1   1]\n",
      " [  0   0   0 ...   4  42   2]\n",
      " [  0   0   0 ...   0   0  67]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eec52583d041432f919784cbe07d6983",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 9:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e1c41cede51f4c918468c9568043bdf1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 9:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  9\n",
      "Accuracy:  0.7580419580419581\n",
      "F1 Score:  0.750557764118338\n",
      "Confusion Matrix:  [[ 43   0   0 ...   1   0   0]\n",
      " [  0  76   0 ...   0   0   0]\n",
      " [  0   0 133 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  45   4   0]\n",
      " [  0   0   0 ...   2  45   1]\n",
      " [  1   0   0 ...   0   1  66]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f622b0eb883e40818f4b33bebd1ccc92",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 10:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9cf41615dab642709bb3126472b37901",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 10:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  10\n",
      "Accuracy:  0.7510489510489511\n",
      "F1 Score:  0.7399803376737685\n",
      "Confusion Matrix:  [[ 41   0   1 ...   0   0   0]\n",
      " [  0  76   0 ...   0   0   0]\n",
      " [  0   0 131 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  41   6   0]\n",
      " [  0   0   0 ...   1  48   1]\n",
      " [  0   0   0 ...   0   1  63]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b37e179abaff42329575fa3d80d0d4e3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 11:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "73c369b887a34e04aa563ae4c8368bd8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 11:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  11\n",
      "Accuracy:  0.7487179487179487\n",
      "F1 Score:  0.7416898027998126\n",
      "Confusion Matrix:  [[ 40   0   0 ...   1   0   0]\n",
      " [  0  77   0 ...   0   0   0]\n",
      " [  0   0 118 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  48   0   0]\n",
      " [  0   0   0 ...   3  45   1]\n",
      " [  0   0   0 ...   1   2  60]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "27955c6a7ac24558bd65fcdfc7c6ab4f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 12:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4e5a76e0d8e4ea881a3040999befc15",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 12:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  12\n",
      "Accuracy:  0.7524475524475525\n",
      "F1 Score:  0.7439707795525562\n",
      "Confusion Matrix:  [[ 33   0   0 ...   1   0   0]\n",
      " [  0  73   0 ...   0   0   0]\n",
      " [  0   0 130 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  41   2   0]\n",
      " [  0   0   0 ...   3  45   1]\n",
      " [  0   0   0 ...   0   2  62]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8b06453cced5455695c8e23d46374f16",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 13:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34cfa7b64b3341ccbc2f8176b9bc2222",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 13:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  13\n",
      "Accuracy:  0.7525641025641026\n",
      "F1 Score:  0.7422151185298976\n",
      "Confusion Matrix:  [[ 26   0   0 ...   0   0   0]\n",
      " [  0  76   0 ...   0   0   0]\n",
      " [  0   0 123 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  45   1   0]\n",
      " [  0   0   0 ...   1  47   0]\n",
      " [  0   0   0 ...   0   2  63]]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1cb420dc74214302afca6cb350cfd9d0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Train Epoch 14:   0%|          | 0/375 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2672509338d645be8bd35e932d3b0893",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Test Epoch 14:   0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:  14\n",
      "Accuracy:  0.7573426573426574\n",
      "F1 Score:  0.7488799597502518\n",
      "Confusion Matrix:  [[ 37   0   0 ...   0   0   1]\n",
      " [  0  74   0 ...   0   0   0]\n",
      " [  0   0 116 ...   0   0   0]\n",
      " ...\n",
      " [  0   0   0 ...  45   0   0]\n",
      " [  0   0   0 ...   1  43   2]\n",
      " [  0   0   0 ...   0   2  65]]\n"
     ]
    }
   ],
   "source": [
    "from torchvision.models import resnet18\n",
    "\n",
    "# Load the pre-trained model\n",
    "model = resnet18(weights=None)\n",
    "model.load_state_dict(torch.load(\"resnet18-f37072fd.pth\"))\n",
    "\n",
    "# Change the output layer\n",
    "model.fc = torch.nn.Linear(512, len(categories))\n",
    "model = model.to(device)\n",
    "\n",
    "# Define the loss function\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=6, gamma=0.1)\n",
    "\n",
    "# Define the data loaders\n",
    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)\n",
    "\n",
    "# Train the model\n",
    "model.train()\n",
    "\n",
    "for epoch in range(15):\n",
    "    for images, labels in tqdm(train_loader, desc=\"Train Epoch \" + str(epoch)):\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    # Evaluate the model\n",
    "    model.eval()\n",
    "    y_true = []\n",
    "    y_pred = []\n",
    "    with torch.no_grad():\n",
    "        for images, labels in tqdm(test_loader, desc=\"Test Epoch \" + str(epoch)):\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            y_true.extend(labels.cpu().numpy())\n",
    "            y_pred.extend(predicted.cpu().numpy())\n",
    "    model.train()\n",
    "    print(\"Epoch: \", epoch)\n",
    "    print(\"Accuracy: \", metrics.accuracy_score(y_true, y_pred))\n",
    "    print(\"F1 Score: \", metrics.f1_score(y_true, y_pred, average='macro'))\n",
    "    print(\"Confusion Matrix: \", metrics.confusion_matrix(y_true, y_pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "79e04ecf5c214316b3aa26cceaa5b666",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/269 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy:  0.7573426573426574\n"
     ]
    }
   ],
   "source": [
    "# Evaluate on the test set\n",
    "\n",
    "model.eval()\n",
    "y_true = []\n",
    "y_pred = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, labels in tqdm(test_loader):\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "        outputs = model(images)\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "        y_true.extend(labels.cpu().numpy())\n",
    "        y_pred.extend(predicted.cpu().numpy())\n",
    "\n",
    "print(\"Accuracy: \", metrics.accuracy_score(y_true, y_pred))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                          precision    recall  f1-score   support\n",
      "\n",
      "                     n02085620-Chihuahua       0.64      0.71      0.67        52\n",
      "              n02085782-Japanese_spaniel       0.82      0.87      0.85        85\n",
      "                   n02085936-Maltese_dog       0.87      0.76      0.81       152\n",
      "                      n02086079-Pekinese       0.73      0.78      0.75        49\n",
      "                      n02086240-Shih-Tzu       0.69      0.72      0.70       114\n",
      "              n02086646-Blenheim_spaniel       0.87      0.86      0.87        88\n",
      "                      n02086910-papillon       0.97      0.91      0.94        96\n",
      "                   n02087046-toy_terrier       0.83      0.67      0.74        72\n",
      "           n02087394-Rhodesian_ridgeback       0.72      0.40      0.52        72\n",
      "                  n02088094-Afghan_hound       0.95      0.95      0.95       139\n",
      "                        n02088238-basset       0.71      0.80      0.75        75\n",
      "                        n02088364-beagle       0.86      0.62      0.72        95\n",
      "                    n02088466-bloodhound       0.93      0.76      0.84        87\n",
      "                      n02088632-bluetick       0.88      0.73      0.80        71\n",
      "       n02089078-black-and-tan_coonhound       0.95      0.61      0.74        59\n",
      "                  n02089867-Walker_hound       0.47      0.66      0.55        53\n",
      "              n02089973-English_foxhound       0.72      0.63      0.67        57\n",
      "                       n02090379-redbone       0.61      0.75      0.67        48\n",
      "                        n02090622-borzoi       0.65      0.78      0.71        51\n",
      "               n02090721-Irish_wolfhound       0.99      0.58      0.73       118\n",
      "             n02091032-Italian_greyhound       0.75      0.67      0.71        82\n",
      "                       n02091134-whippet       0.73      0.69      0.71        87\n",
      "                  n02091244-Ibizan_hound       0.92      0.82      0.87        88\n",
      "            n02091467-Norwegian_elkhound       0.95      0.91      0.93        96\n",
      "                    n02091635-otterhound       0.71      0.92      0.80        51\n",
      "                        n02091831-Saluki       0.84      0.80      0.82       100\n",
      "            n02092002-Scottish_deerhound       0.76      0.86      0.81       132\n",
      "                    n02092339-Weimaraner       0.96      0.82      0.88        60\n",
      "     n02093256-Staffordshire_bullterrier       0.72      0.47      0.57        55\n",
      "n02093428-American_Staffordshire_terrier       0.54      0.69      0.61        64\n",
      "            n02093647-Bedlington_terrier       0.94      0.90      0.92        82\n",
      "                n02093754-Border_terrier       0.89      0.93      0.91        72\n",
      "            n02093859-Kerry_blue_terrier       0.86      0.71      0.78        79\n",
      "                 n02093991-Irish_terrier       0.88      0.75      0.81        69\n",
      "               n02094114-Norfolk_terrier       0.69      0.74      0.71        72\n",
      "               n02094258-Norwich_terrier       0.73      0.73      0.73        85\n",
      "             n02094433-Yorkshire_terrier       0.49      0.72      0.59        64\n",
      "       n02095314-wire-haired_fox_terrier       0.73      0.63      0.68        57\n",
      "              n02095570-Lakeland_terrier       0.74      0.58      0.65        97\n",
      "              n02095889-Sealyham_terrier       0.98      0.86      0.92       102\n",
      "                      n02096051-Airedale       0.81      0.85      0.83       102\n",
      "                         n02096177-cairn       0.79      0.66      0.72        97\n",
      "            n02096294-Australian_terrier       0.71      0.71      0.71        96\n",
      "                n02096437-Dandie_Dinmont       0.97      0.74      0.84        80\n",
      "                   n02096585-Boston_bull       0.90      0.87      0.88        82\n",
      "           n02097047-miniature_schnauzer       0.68      0.74      0.71        54\n",
      "               n02097130-giant_schnauzer       0.51      0.72      0.60        57\n",
      "            n02097209-standard_schnauzer       0.60      0.47      0.53        55\n",
      "                n02097298-Scotch_terrier       0.70      0.86      0.78        58\n",
      "               n02097474-Tibetan_terrier       0.61      0.75      0.67       106\n",
      "                 n02097658-silky_terrier       0.81      0.52      0.63        83\n",
      "   n02098105-soft-coated_wheaten_terrier       0.73      0.57      0.64        56\n",
      "   n02098286-West_Highland_white_terrier       0.70      0.86      0.77        69\n",
      "                         n02098413-Lhasa       0.54      0.62      0.57        86\n",
      "         n02099267-flat-coated_retriever       0.75      0.79      0.77        52\n",
      "        n02099429-curly-coated_retriever       0.81      0.67      0.73        51\n",
      "              n02099601-golden_retriever       0.78      0.78      0.78        50\n",
      "            n02099712-Labrador_retriever       0.58      0.75      0.65        71\n",
      "      n02099849-Chesapeake_Bay_retriever       0.61      0.90      0.72        67\n",
      "   n02100236-German_short-haired_pointer       0.80      0.85      0.82        52\n",
      "                        n02100583-vizsla       0.83      0.80      0.81        54\n",
      "                n02100735-English_setter       0.72      0.70      0.71        61\n",
      "                  n02100877-Irish_setter       0.87      0.82      0.84        55\n",
      "                 n02101006-Gordon_setter       0.87      0.91      0.89        53\n",
      "              n02101388-Brittany_spaniel       0.85      0.75      0.80        52\n",
      "                       n02101556-clumber       0.83      0.86      0.84        50\n",
      "              n02102040-English_springer       0.81      0.92      0.86        59\n",
      "        n02102177-Welsh_springer_spaniel       0.90      0.76      0.83        50\n",
      "                n02102318-cocker_spaniel       0.81      0.71      0.76        59\n",
      "                n02102480-Sussex_spaniel       0.92      0.88      0.90        51\n",
      "           n02102973-Irish_water_spaniel       0.74      0.84      0.79        50\n",
      "                        n02104029-kuvasz       0.48      0.74      0.58        50\n",
      "                    n02104365-schipperke       0.82      0.83      0.83        54\n",
      "                   n02105056-groenendael       0.82      0.92      0.87        50\n",
      "                      n02105162-malinois       0.75      0.78      0.76        50\n",
      "                        n02105251-briard       0.78      0.69      0.73        52\n",
      "                        n02105412-kelpie       0.57      0.81      0.67        53\n",
      "                      n02105505-komondor       0.92      0.89      0.91        54\n",
      "          n02105641-Old_English_sheepdog       0.82      0.81      0.82        69\n",
      "             n02105855-Shetland_sheepdog       0.93      0.49      0.64        57\n",
      "                        n02106030-collie       0.55      0.68      0.61        53\n",
      "                 n02106166-Border_collie       0.69      0.74      0.71        50\n",
      "          n02106382-Bouvier_des_Flandres       0.70      0.64      0.67        50\n",
      "                    n02106550-Rottweiler       0.67      0.98      0.80        52\n",
      "               n02106662-German_shepherd       0.75      0.85      0.79        52\n",
      "                      n02107142-Doberman       0.75      0.72      0.73        50\n",
      "            n02107312-miniature_pinscher       0.85      0.71      0.77        84\n",
      "    n02107574-Greater_Swiss_Mountain_dog       0.59      0.79      0.68        68\n",
      "          n02107683-Bernese_mountain_dog       0.92      0.92      0.92       118\n",
      "                   n02107908-Appenzeller       0.49      0.65      0.56        51\n",
      "                   n02108000-EntleBucher       0.90      0.62      0.73       102\n",
      "                         n02108089-boxer       0.60      0.67      0.63        51\n",
      "                  n02108422-bull_mastiff       0.78      0.77      0.77        56\n",
      "               n02108551-Tibetan_mastiff       0.62      0.77      0.69        52\n",
      "                n02108915-French_bulldog       0.80      0.86      0.83        59\n",
      "                    n02109047-Great_Dane       0.63      0.64      0.64        56\n",
      "                 n02109525-Saint_Bernard       0.93      0.96      0.94        70\n",
      "                    n02109961-Eskimo_dog       0.30      0.58      0.39        50\n",
      "                      n02110063-malamute       0.56      0.69      0.62        78\n",
      "                n02110185-Siberian_husky       0.58      0.15      0.24        92\n",
      "                 n02110627-affenpinscher       0.87      0.78      0.82        50\n",
      "                       n02110806-basenji       0.89      0.79      0.83       109\n",
      "                           n02110958-pug       0.87      0.84      0.85       100\n",
      "                      n02111129-Leonberg       0.89      0.91      0.90       110\n",
      "                  n02111277-Newfoundland       0.65      0.77      0.71        95\n",
      "                n02111500-Great_Pyrenees       0.73      0.58      0.64       113\n",
      "                       n02111889-Samoyed       0.69      0.98      0.81       118\n",
      "                    n02112018-Pomeranian       0.98      0.79      0.87       119\n",
      "                          n02112137-chow       0.87      0.99      0.93        96\n",
      "                      n02112350-keeshond       0.90      0.95      0.92        58\n",
      "             n02112706-Brabancon_griffon       0.93      0.81      0.87        53\n",
      "                      n02113023-Pembroke       0.84      0.77      0.80        81\n",
      "                      n02113186-Cardigan       0.67      0.75      0.71        55\n",
      "                    n02113624-toy_poodle       0.58      0.43      0.49        51\n",
      "              n02113712-miniature_poodle       0.35      0.53      0.42        55\n",
      "               n02113799-standard_poodle       0.60      0.66      0.63        59\n",
      "              n02113978-Mexican_hairless       0.92      0.82      0.87        55\n",
      "                         n02115641-dingo       0.67      0.80      0.73        56\n",
      "                         n02115913-dhole       0.96      0.86      0.91        50\n",
      "           n02116738-African_hunting_dog       0.80      0.94      0.87        69\n",
      "\n",
      "                                accuracy                           0.76      8580\n",
      "                               macro avg       0.76      0.75      0.75      8580\n",
      "                            weighted avg       0.78      0.76      0.76      8580\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(metrics.classification_report(y_true, y_pred, target_names=list(categories.keys())))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
